package middlex

import (
	"net/http"

	"github.com/tal-tech/go-zero/core/stores/cache"

	"gitee.com/yookie316/goutils/errorx"

	"github.com/tal-tech/go-zero/core/limit"
	"github.com/tal-tech/go-zero/core/logx"
	"github.com/tal-tech/go-zero/core/stores/redis"
	"github.com/tal-tech/go-zero/rest/httpx"
)

const (
	seconds = 1
	// quota     = 10
	limPrefix = "pl"
)

type RateMiddleware struct {
	lim *limit.PeriodLimit
}

func NewRateMiddleware(conf cache.CacheConf, totalPerSecond int) *RateMiddleware {
	redisConf := conf[0]
	store := redis.NewRedis(redisConf.Host, redisConf.Type, redisConf.Pass)
	l := limit.NewPeriodLimit(seconds, totalPerSecond, store, limPrefix)
	return &RateMiddleware{
		lim: l,
	}
}

func (m *RateMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		code, err := m.lim.Take("first")
		if err != nil {
			logx.Errorf("get lim error:%s", err)
			httpx.Error(w, errorx.NewInternetError())
			return
		}
		pass := false
		switch code {
		case limit.OverQuota:
			logx.Errorf("OverQuota key: %v", limPrefix)
		case limit.Allowed:
			pass = true
		case limit.HitQuota:
			logx.Errorf("HitQuota key: %v", limPrefix)
			// todo: maybe we need to let users know they hit the quota
		default:
			logx.Errorf("DefaultQuota key: %v", limPrefix)
			// unknown response, we just let the sms go
			pass = true
		}
		if !pass {
			httpx.Error(w, errorx.NewTooManyRequestError())
			return
		}
		next(w, r)
	}
}
